!--------------------------------------------------------------------------------------------------!
!  DFTB+: general package for performing fast atomistic simulations                                !
!  Copyright (C) 2006 - 2020  DFTB+ developers group                                               !
!                                                                                                  !
!  See the LICENSE file for terms of usage and distribution.                                       !
!--------------------------------------------------------------------------------------------------!

#:include 'common.fypp'

!> Program for plotting molecular orbitals as cube files.
program waveplot
  use dftbp_assert
  use dftbp_globalenv, only : stdOut
  use dftbp_initwaveplot
  use dftbp_accuracy
  use dftbp_charmanip
  use dftbp_fileid
  use dftbp_typegeometry
  use dftbp_fileid
  use dftbp_gridcache
  use dftbp_molecularorbital
  use dftbp_simplealgebra
  use dftbp_linkedlist
  use dftbp_periodic
#:if WITH_MPI
  use dftbp_mpienv, only : TMpiEnv, TMpiEnv_init
  use dftbp_mpifx, only : mpifx_init_thread, mpifx_finalize
  use mpi, only : MPI_THREAD_FUNNELED
#:endif
  implicit none

  character(len=80) :: comments(2), fileName
  real(dp), pointer :: gridValReal(:,:,:)
  complex(dp), pointer :: gridValCmpl(:,:,:)
  real(dp), allocatable :: buffer(:,:,:), totChrg(:,:,:), atomicChrg(:,:,:,:)
  real(dp), allocatable :: spinUp(:,:,:)
  real(dp), allocatable :: orbitalOcc(:,:)
  real(dp), allocatable :: coords(:,:)
  real(dp) :: shift(3)
  integer, allocatable :: species(:)
  integer :: nBox
  real(dp) :: sumTotChrg, sumChrg, sumAtomicChrg
  integer :: levelIndex(3)
  integer :: iLevel, iKPoint, iSpin, iSpinOld
  integer :: iAtom, iSpecies, iAng, mAng, ind
  logical :: tFinished, tPlotLevel
  real(dp) :: invBoxVecs(3,3), recVecs2p(3,3)
  real(dp) :: cellMiddle(3), boxMiddle(3), frac(3), cubeCorner(3), coord(3)
  real(dp) :: mDist, dist
  real(dp), allocatable :: cellVec(:,:), rCellVec(:,:)
  integer :: i1, i2, i3, iCell
  type(TListRealR1) :: coordList
  type(TListInt) :: speciesList

#:if WITH_MPI
  !> MPI environment, if compiled with mpifort
  type(TMpiEnv) :: mpiEnv

  ! As this is serial code, trap for run time execution on more than 1 processor with an mpi enabled
  ! build
  call mpifx_init_thread(requiredThreading=MPI_THREAD_FUNNELED)
  call TMpiEnv_init(mpiEnv)
  call mpiEnv%mpiSerialEnv()
#:endif

  ! Allocate resources
  call initProgramVariables()
  write(stdout, "(/,A,/)") "Starting main program"

  ! Allocating buffer for general grid, total charge and spin up
  allocate(buffer(nPoints(1), nPoints(2), nPoints(3)))
  if (tCalcTotChrg) then
    allocate(totChrg(nPoints(1), nPoints(2), nPoints(3)))
    totChrg = 0.0_dp
    if (tPlotTotSpin) then
      allocate(spinUp(nPoints(1), nPoints(2), nPoints(3)))
    end if
  end if
  write (comments(1), "(A)") "Cube file generated by WAVEPLOT from data &
      &created by DFTB+"

  ! Repeat boxes if necessary
  nBox = product(repeatBox)
  if (nBox > 1) then
    ! If tFillBox is off, coordinates must be repeated here. Otherwise the part for filling with
    ! atoms will do that.
    if (.not. tFillBox) then
      allocate(coords(3, size(geo%coords)))
      allocate(species(size(geo%species)))
      coords(:,:) = geo%coords(:,:)
      species(:) = geo%species(:)
      deallocate(geo%coords)
      deallocate(geo%species)
      allocate(geo%coords(3, nBox * geo%nAtom))
      allocate(geo%species(nBox * geo%nAtom))
      ind = 0
      do i1 = 0, repeatBox(1) - 1
        do i2 = 0, repeatBox(2) - 1
          do i3 = 0, repeatBox(3) - 1
            shift(:) = matmul(boxVecs, real((/ i1, i2, i3/), dp))
            do iAtom = 1, geo%nAtom
              geo%coords(:,ind+iAtom) = coords(:,iAtom) + shift(:)
            end do
            geo%species(ind+1:ind+geo%nAtom) = species(:)
            ind = ind + geo%nAtom
          end do
        end do
      end do
      geo%nAtom = nBox * geo%nAtom
    end if
    do i1 = 1, 3
      boxVecs(:,i1) = boxVecs(:,i1) * real(repeatBox(i1), dp)
    end do
  end if

  write(stdout, "(A)") "Origin"
  write(stdout, "(2X,3(F0.5,1X))") origin(:)
  write(stdout, "(A)") "Box"
  do i1 = 1, 3
    write(stdout, "(2X,3(F0.5,1X))") boxVecs(:, i1)
  end do
  write(stdout, "(A)") "Spatial resolution [1/Bohr]:"
  write(stdout, "(2X,3(F0.5,1X))") 1.0_dp / sqrt(sum(gridVec**2, dim=1))
  write(stdout, *)

  ! Create density superposition of the atomic orbitals. Occupation is distributed equally on
  ! orbitals with the same angular momentum.
  if (tCalcAtomDens) then
    allocate(atomicChrg(nPoints(1), nPoints(2), nPoints(3), 1))
    allocate(orbitalOcc(nOrb, 1))
    ind = 1
    do iAtom = 1, geo%nAtom
      iSpecies = geo%species(iAtom)
      do iAng = 1, size(basis(iSpecies)%angMoms)
        mAng = 2 * basis(iSpecies)%angMoms(iAng) + 1
        orbitalOcc(ind:ind+mAng-1,1) = basis(iSpecies)%occupations(iAng) &
            &/ real(mAng, dp)
        ind = ind + mAng
      end do
    end do
    call getValue(molorb, gridOrigin, gridVec, orbitalOcc, atomicChrg, &
        &addDensities=.true.)
    sumAtomicChrg = sum(atomicChrg) * gridVol

    if (tVerbose) then
      write(stdout, "('Total charge of atomic densities:',F12.6,/)") sumAtomicChrg
    end if
    if (tPlotAtomDens) then
      write (comments(2), 9989) identity
9989  format('Calc-Id:',I11,', atomdens')
      fileName = "wp-atomdens.cube"
      call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
          &fileName, comments, repeatBox)
      write(stdout, "(A)") "File '" // trim(fileName) // "' written"
    end if
  end if

  if (tVerbose) then
    write(stdout, "(/,A5,' ',A6,' ',A6,' ',A7,' ',A11,' ',A11)") "Spin", "KPoint", &
        &"State", "Action", "Norm", "W. Occup."
  end if

  ! Fold in coordinates
  if (tFoldCoords) then
    call invert33(invBoxVecs, boxVecs)
    call invert33(recVecs2p, geo%latVecs)
    recVecs2p = reshape(recVecs2p, (/3, 3/), order=(/2, 1/))
    call foldCoordToUnitCell(geo%coords(:,:), geo%latVecs, recVecs2p)
  end if

  ! Fill the box with atoms
  if (tFillBox) then
    ! Shifting plotted region by integer lattice vectors, to have its center as close to the center
    ! of the lattice unit cell as possible.
    cellMiddle(:) = 0.5_dp * sum(geo%latVecs, dim=2)
    boxMiddle(:) = origin(:) + 0.5_dp * sum(boxVecs, dim=2)
    frac(:) = matmul(boxMiddle - cellMiddle, recVecs2p)
    origin(:) = origin(:) - matmul(geo%latVecs, real(anint(frac), dp))
    gridOrigin(:) = gridOrigin(:) - matmul(geo%latVecs, real(anint(frac), dp))
    ! We need all cells around, which could contain atoms in the sphere, drawn from the center of
    ! the unit cell, containing entire plotted region
    mDist = 0.0_dp
    do i1 = 0, 1
      do i2 = 0, 1
        do i3 = 0, 1
          cubeCorner(:) = origin(:) &
              &+ matmul(boxVecs, real((/ i1, i2, i3 /), dp))
          dist = sqrt(sum((cubeCorner - cellMiddle)**2))
          mDist = max(dist, mDist)
        end do
      end do
    end do
    call getCellTranslations(cellVec, rCellVec, geo%latVecs, recVecs2p, mDist)
    ! Check all atoms in the shifted cells, if they fall in the plotted region
    call init(coordList)
    call init(speciesList)
    do iCell = 1, size(rCellVec, dim=2)
      do iAtom = 1, geo%nAtom
        coord(:) = geo%coords(:,iAtom) + rCellVec(:,iCell)
        frac(:) = matmul(invBoxVecs, coord - origin)
        if (all(frac > -1e-4_dp) .and. all(frac < 1.0_dp + 1e-4_dp)) then
          call append(coordList, coord)
          call append(speciesList, geo%species(iAtom))
        end if
      end do
    end do
    deallocate(geo%coords)
    deallocate(geo%species)
    geo%nAtom = len(coordList)
    allocate(geo%coords(3, geo%nAtom))
    allocate(geo%species(geo%nAtom))
    call asArray(coordList, geo%coords)
    call asArray(speciesList, geo%species)
    deallocate(cellVec)
    deallocate(rCellVec)
  end if

  ! Calculate the molecular orbitals and write them to the disc
  iSpinOld = 1
  tFinished = .false.
  lpStates: do while (.not. tFinished)
    ! Get the next grid and its parameters
    if (tRealHam) then
      call next(grid, gridValReal, levelIndex, tFinished)
    else
      call next(grid, gridValCmpl, levelIndex, tFinished)
    end if
    iLevel = levelIndex(1)
    iKPoint = levelIndex(2)
    iSpin = levelIndex(3)

    ! Save spin up density before processing first level for spin down
    if (tPlotTotSpin .and. (iSpinOld /= iSpin)) then
      iSpinOld = iSpin
      spinUp(:,:,:) = totChrg(:,:,:)
    end if

    ! Build charge if needed for total charge or was explicitely required
    tPlotLevel = any(plottedSpins == iSpin) &
        &.and. any(plottedKPoints == iKPoint) .and. any(plottedLevels == iLevel)
    if (tCalcTotChrg .or. (tPlotLevel .and. (tPlotChrg.or.tPlotChrgDiff))) then
      if (tRealHam) then
        buffer = gridValReal**2
      else
        buffer = abs(gridValCmpl)**2
      end if
      if (tCalcTotChrg) then
        totChrg(:,:,:) = totChrg(:,:,:) &
            &+ occupations(iLevel, iKPoint, iSpin) * buffer(:,:,:)
      end if
      sumChrg = sum(buffer) * gridVol
      if (tVerbose) then
        write(stdout, "(I5,I7,I7,A8,F12.6,F12.6)") iSpin, iKPoint, iLevel, "calc",&
            &sumChrg, occupations(iLevel, iKPoint, iSpin)
      end if
    end if

    ! Build and dump desired properties of the current level
    if (tPlotLevel) then
      if (tPlotChrg) then
        write (comments(2), 9990) identity, iSpin, iKPoint, iLevel
9990    format('Calc-Id:',I11,', Spin:',I2,', K-Point:',I6,', State:',I6, &
            &', abs2')
        fileName = "wp-" // i2c(iSpin) // "-" // i2c(iKPoint) // "-" &
            &//i2c(iLevel) // "-abs2.cube"
        call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
            &fileName, comments, repeatBox)
        write(stdout, "(A)") "File '" // trim(fileName) // "' written"
      end if

      if (tPlotChrgDiff) then
        buffer = buffer - (sumChrg / sumAtomicChrg) * atomicChrg(:,:,:,1)
        write (comments(2), 9995) identity, iSpin, iKPoint, iLevel
9995    format('Calc-Id:',I11,', Spin:',I2,', K-Point:',I6,', State:',I6, &
            &', abs2diff')
        fileName = "wp-" // i2c(iSpin) // "-" // i2c(iKPoint) // "-" &
            &//i2c(iLevel) // "-abs2diff.cube"
        call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
            &fileName, comments, repeatBox)
        write(stdout, "(A)") "File '" // trim(fileName) // "' written"
      end if

      if (tPlotReal) then
        if (tRealHam) then
          buffer = gridValReal
        else
          buffer = real(gridValCmpl, dp)
        end if
        write (comments(2), 9991) identity, iSpin, iKPoint, iLevel
9991    format('Calc-Id:',I11,', Spin:',I2,', K-Point:',I6,', State:',I6, &
            &', real')
        fileName = "wp-" // i2c(iSpin) // "-" // i2c(iKPoint) // "-" &
            &//i2c(iLevel) // "-real.cube"
        call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
            &fileName, comments, repeatBox)
        write(stdout, "(A)") "File '" // trim(fileName) // "' written"
      end if

      if (tPlotImag) then
        buffer = aimag(gridValCmpl)
        write (comments(2), 9992) identity, iSpin, iKPoint, iLevel
9992    format('Calc-Id:',I11,', Spin:',I2,', K-Point:',I6,', State:',I6, &
            &', imag')
        fileName = "wp-" // i2c(iSpin) // "-" // i2c(iKPoint) // "-" &
            &//i2c(iLevel) // "-imag.cube"
        call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
            &fileName, comments, repeatBox)
        write(stdout, "(A)") "File '" // trim(fileName) // "' written"
      end if
    end if
  end do lpStates

  ! Dump total charge, if required
  if (tCalcTotChrg) then
    sumTotChrg = sum(totChrg) * gridVol
  end if
  if (tPlotTotChrg) then
    write (comments(2), 9993) identity
9993 format('Calc-Id:',I11,', abs2')
    fileName = "wp-abs2.cube"
    call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, totChrg, &
        &fileName, comments, repeatBox)
    write(stdout, "(A)") "File '" // trim(fileName) // "' written"
    if (tVerbose) then
      write(stdout, "(/,'Total charge:',F12.6,/)") sumTotChrg
    end if
  end if

  ! Dump total charge difference
  if (tPlotTotDiff) then
    buffer = totChrg - (sumTotChrg / sumAtomicChrg) * atomicChrg(:,:,:,1)
    write (comments(2), 9994) identity
9994 format('Calc-Id:',I11,', abs2diff')
    fileName = 'wp-abs2diff.cube'
    call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
        &fileName, comments, repeatBox)
    write(stdout, "(A)") "File '" // trim(fileName) // "' written"
  end if

  ! Dump spin polarisation
  if (tPlotTotSpin) then
    buffer = 2.0_dp * spinUp - totChrg
    write (comments(2), 9996) identity
9996 format('Calc-Id:',I11,', spinpol')
    fileName = 'wp-spinpol.cube'
    call writeCubeFile(geo, atomicNumbers, gridVec, gridOrigin, buffer, &
        &fileName, comments, repeatBox)
    write(stdout, "(A)") "File '" // trim(fileName) // "' written"
  end if

#:if WITH_MPI
  call mpifx_finalize()
#:endif


contains


  !> Writes a 3D function as cube file.
  subroutine writeCubeFile(geo, atomicNumbers, gridVecs, origin, gridVal, &
      &fileName, comments, repeatBox)

    !> Geometry information about the structure
    type(TGeometry), intent(in) :: geo

    !> Atomic numbers correpsonding to each species
    integer, intent(in) :: atomicNumbers(:)

    !> Grid vectors
    real(dp), intent(in) :: gridVecs(:,:)

    !> Origin of the grid
    real(dp), intent(in) :: origin(:)

    !> Value of the 3D function in the grid points
    real(dp), intent(in) :: gridVal(:,:,:)

    !> Name of the file to create
    character(len=*), intent(in) :: fileName

    !> First two comment lines of the file
    character(len=*), intent(in), optional :: comments(:)

    !> How often the grid should be repeated along the direction of the grid vectors
    integer, intent(in), optional :: repeatBox(:)

    integer, parameter :: bufferSize = 6
    real(dp) :: buffer(bufferSize)
    character(len=*), parameter :: formBuffer = "(6E16.8)"
    integer :: rep(3)
    integer, save :: fd = -1
    integer :: ii, i1, i2, i3, ir1, ir2, ir3

    @:ASSERT(size(atomicNumbers) == size(geo%speciesNames))
    @:ASSERT(all(shape(gridVecs) == (/3, 3/)))
    @:ASSERT(size(origin) == 3)
    @:ASSERT(all(shape(gridVal) >= (/ 0, 0, 0 /)))
  #:block DEBUG_CODE
    if (present(comments)) then
      @:ASSERT(size(comments) == 2)
    end if
    if (present(repeatBox)) then
      @:ASSERT(size(repeatBox) == 3)
      @:ASSERT(all(repeatBox > 0))
    end if
  #:endblock DEBUG_CODE

    if (present(repeatBox)) then
      rep(:) = repeatBox(:)
    else
      rep(:) = (/ 1, 1, 1 /)
    end if

    if (fd == -1) then
      fd = getFileId()
    end if

    open(fd, file=fileName, action="write", status="replace")
    if (present(comments)) then
      write (fd, "(A)") trim(comments(1))
      write (fd, "(A)") trim(comments(2))
    else
      write (fd, "(A)") "Made by waveplot"
      write (fd, *)
    end if
    write (fd,"(I5,3F12.6)") geo%nAtom, origin(:)
    write (fd,"(I5,3F12.6)") rep(1) * size(gridVal, dim=1), gridVecs(:,1)
    write (fd,"(I5,3F12.6)") rep(2) * size(gridVal, dim=2), gridVecs(:,2)
    write (fd,"(I5,3F12.6)") rep(3) * size(gridVal, dim=3), gridVecs(:,3)
    do ii = 1, geo%nAtom
      write (fd, "(I5,4F12.6)") atomicNumbers(geo%species(ii)), 0.0_dp, &
          &geo%coords(:, ii)
    end do

    do ir1 = 1, rep(1)
      do i1 = 1, size(gridVal, dim=1)
        do ir2 = 1, rep(2)
          do i2 = 1, size(gridVal, dim=2)
            do ir3 = 1, rep(3)
              do i3 = 1, size(gridVal, dim=3)
                ii = mod(i3-1, bufferSize) + 1
                buffer(ii) = gridVal(i1, i2, i3)
                if (ii == bufferSize) then
                  write (fd,formBuffer) real(buffer)
                end if
              end do
              if (ii /= bufferSize) then
                write (fd, "(" // i2c(ii) // "E16.8)") real(buffer(:ii))
              end if
            end do
          end do
        end do
      end do
    end do
    close(fd)

  end subroutine writeCubeFile

end program waveplot
